from sklearn.svm import SVC
from sklearn.metrics import classification_report
from text_processor import TextProcessor

class EnhancedClassifier:
    def __init__(self, vector_size=100, window=5, min_count=1):
        self.text_processor = TextProcessor(vector_size, window, min_count)
        self.classifier = SVC(kernel='linear')
        self.is_trained = False
        
    def train(self, documents, labels):
        """训练分类器"""
        # 处理所有文档，生成特征向量
        X = self.text_processor.process_documents(documents)
        y = labels
        
        # 训练SVM分类器
        self.classifier.fit(X, y)
        self.is_trained = True
        
    def predict(self, document):
        """预测文档的类别"""
        if not self.is_trained:
            raise ValueError("分类器尚未训练")
            
        # 生成文档的特征向量
        feature_vector = self.text_processor.get_document_vector(document)
        
        # 预测类别
        return self.classifier.predict([feature_vector])[0]
    
    def evaluate(self, test_documents, test_labels):
        """评估分类器性能"""
        if not self.is_trained:
            raise ValueError("分类器尚未训练")
            
        # 生成测试文档的特征向量
        X_test = [self.text_processor.get_document_vector(doc) for doc in test_documents]
        
        # 进行预测
        predictions = self.classifier.predict(X_test)
        
        # 生成评估报告
        return classification_report(test_labels, predictions)